#include "CoordinateProjection.h"
#include "Constraint.h"
#include "MatrixClasses.h"
#include "MultibodySystem.h" 
#include "Utils.h"

#include <iomanip>

using namespace RSIM;
using namespace std;

/////////////////////////////////////////////////////// 

CoordinateProjection::~CoordinateProjection(){        
        delete [] CJac_;
        delete [] Cq_;
        delete [] delq_;
}

/////////////////////////////////////////////////////// 

void CoordinateProjection::setConstraintTolerance(const double& Tol){
        CTol_ = Tol;
}

/////////////////////////////////////////////////////// 

void CoordinateProjection::setMaxIterations(const int& niter){
        MaxCprIter_ = niter;
}

///////////////////////////////////////////////////////

void CoordinateProjection::calcPositionLevelErrors(){                
        
        // calculate the position level errors for all constraints. 
        // This also updates the cumulative beta vector if applicable.
        this->calcPositionLevelErrorForAllConstraints(constraints_); 
        
        // This copies the position level errors computed in previous step into the Cq Vector
        // and makes it ready for coordinate projection.
        this->popuplateCqVector(constraints_,ccouple_,Cq_);
        
        // This calculates the norms of error at position level in the
        // coupled constraints which are to be satisfied simultaneously.
        this->calcPositionErrorNorms(Cq_,Nccouple_,PositionErrorNorms_);
}

///////////////////////////////////////////////////////

void CoordinateProjection::attemptSatisfyPositionLevelConstraint(){
        
        int num_iter = 0;
        while(1){                                
                Utils::rsim_assert(num_iter++<MaxCprIter_,
                                   "Unable to meet constraint tolerance. Either reduce the constraint tolerance or increase the number of coordinate projection iterations." );
                bool success = true;
                
                // First check for position level errors. 
                calcPositionLevelErrors();
                        
                for(int i=0;i<Nccouple_;++i){
                        // Check if position error is less than tolerance
                        bool isErrorNormAcceptable = PositionErrorNorms_[i] < CTol_;
                        //cout<<setprecision(16);
                        cout<<"Iteration Number = "<<num_iter<<" Error in Position Norm = "<<PositionErrorNorms_[i]<<endl;
                        
                        // success is true only if the error norm in all the coupled 
                        // constraints is below a certain threshold.
                        success = success && isErrorNormAcceptable;
                        
                        if(isErrorNormAcceptable){                        
                                // If the error is less than tolerance, nothing else needs to be done.
                                continue;
                        }
                        else{                        
                                // Calculate constraint jacobian for the i^th set of coupled constraints
                                calcOneConstraintJacobian(i);
                                                                
                                // This returns the solution to the problem, minimize 
                                // dq^T dq subject to the constraint CJac*dq = -Cq. Equivalent
                                // to Newton iteration. The coordinates are not scaled in this
                                // version. This has to be done in future.
                                this->calcDelQ(CJac_+i,Cq_+i,delq_+i); // sFlag
                                
                                this->updCoupledConstraintQs(delq_[i],ccouple_[i],constraints_);
//                                 cout<<delq_[i]<<endl;
//                                 cin.get();
                                
                                updSystemPosition();
                        }
                }//for(int i=0;i<Nccouple_;++i)
                
                if(success)
                        break;
        }//while(1)        
}

///////////////////////////////////////////////////////

void CoordinateProjection::calcConstraintJacobians_Tobedeleted(){        
        
        // nJac is the total number of constraint jacobians
        for(int i=0;i<Nccouple_;++i){
                
                // List of coupled constraints which go in i^th constraint jacobian                
                const List <int>& coup_cons = ccouple_[i];
                
                // This loop calculates the i^th constraint Jacobian
                for(int j=0;j<coup_cons.getLength();++j){
                        
                        const int& constraint_id = coup_cons[j];
                        
                        // All coupled constraints contibute to the same constraint jacobian
                        constraints_[constraint_id]->calcConstraintJacobian(ConstraintToUMap_[constraint_id], CJac_+i);
                }
        }
}

///////////////////////////////////////////////////////

void CoordinateProjection::calcOneConstraintJacobian(const int& index){
        // List of coupled constraints which go in i^th constraint jacobian                
        const List <int>& coup_cons = ccouple_[index];
        
        // This loop calculates the i^th constraint Jacobian
        for(int j=0;j<coup_cons.getLength();++j){
                
                const int& constraint_id = coup_cons[j];
                
                // All coupled constraints contibute to the same constraint jacobian
                constraints_[constraint_id]->calcConstraintJacobian(ConstraintToUMap_[constraint_id], CJac_+index);
        }
}

///////////////////////////////////////////////////////

void CoordinateProjection::updSystemPosition(){
        for(int i=0;i<system_.getLength();++i){
                // Create new labels for MultibodySystem and State
                MultibodySystem& MBS = *(system_[i]);
                State& state = *(state_[i]);
                
                MBS.updSystemPosition(state);
        }                
}

///////////////////////////////////////////////////////

void CoordinateProjection::init(){        
        
        // Set the size of Constaints to U mapping.
        ConstraintToUMap_.resize(constraints_.getLength());        
        
        // This differs form the actual number of constraints.
        // nConstraints is number of rows of ccouple_ or includes
        // the information for the constraints which must be solved
        // simultaneously.
        Nccouple_ = ccouple_.getLength();
                                        
        PositionErrorNorms_.resize(Nccouple_);
        VelocityErrorNorms_.resize(Nccouple_);
        
        // Constraint Jacobians which are formulated separately
        CJac_ = new Matrix[Nccouple_];
        Cq_ = new Vector[Nccouple_];
        delq_ = new Vector[Nccouple_];
        
        // This sets the variable sid_
        this->sortSysId(ccouple_,constraints_,sid_);                
        int nu, m;
        
        List <List <int> > bmap;
        List <List <int> > umap_master;
        
        for(int i=0;i<Nccouple_;++i){
                // Get a set of coupled constraints which must be solved simultaneously
                const List <int>& coup_cons = ccouple_[i];
                const int& n_coup_cons = coup_cons.getLength();
                
                this->resize_map(sid_[i].getLength(),bmap);
                this->resize_map(sid_[i].getLength(),umap_master);
                nu = 0;
                m = 0;                        
                for(int j=0;j<n_coup_cons;++j){
                                                        
                        const int& constraint_id = coup_cons[j];
                        Constraint *ptr = constraints_[constraint_id];
                        
                        ConstraintToUMap_[constraint_id] = m;
                                                                                        
                        // Get the umaps for this constraint
                        List <int>& pumap = ptr->wgetPUMap();
                        List <int>& cumap = ptr->wgetCUMap();
                        
                        // Calculate the number of algebraic constraints in the system
                        m += ptr->getNAlgebraicConstraints();
                        
                        const List <int>& p_path =  ptr->getParentPathList();
                        const List <int>& c_path = ptr->getChildPathList();
                        
                        // Cumulative betak is required only if the linear velocity
                        // is constrained. For constrained angular velocity, the partials
                        // are trivially obtained.
                        if(ptr->imposeConstraintOnLinearVelocity()){
                                ptr->allocateSpaceForBetak();                                        
                        }
                        
                        // Get Parent and Child SysId
                        const int& pid = ptr->getParentSysId();
                        const int& cid = ptr->getChildSysId();        
                        
                        int p_sys_index = sid_[i].findElement(pid);
                        int c_sys_index = sid_[i].findElement(cid);                                                                
                        
                        // Loop through parent
                        for(int iloop=0;iloop<p_path.getLength()-1;++iloop){
                                const int& body_index = p_path[iloop];                                        
                                
                                // Check if the body for this particular system is already added
                                int u_index_tmp = bmap[p_sys_index].findElement(body_index);
                                if(!(u_index_tmp+1)){
                                        pumap.addNode(nu);
                                        
                                        bmap[p_sys_index].addNode(body_index);
                                        umap_master[p_sys_index].addNode(nu);
                                        
                                        const RigidBodyNode *Nd = system_[pid]->getNode(body_index);                                                
                                        nu += Nd->getDOF();                                                                                                
                                }
                                else{                                                
                                        pumap.addNode(umap_master[p_sys_index][u_index_tmp]);
                                }
                        } // Parent constrain for loop ends here
                        pumap.finalize();
                        
                        // Loop through child
                        for(int iloop=0;iloop<c_path.getLength()-1;++iloop){
                                const int& body_index = c_path[iloop];                                        
                                
                                // Check if the body for this particular system is already added
                                int u_index_tmp = bmap[c_sys_index].findElement(body_index);                                        
                                if(!(u_index_tmp+1)){
                                        cumap.addNode(nu);
                                        
                                        bmap[c_sys_index].addNode(body_index);
                                        umap_master[c_sys_index].addNode(nu);
                                        
                                        const RigidBodyNode *Nd = system_[cid]->getNode(body_index);
                                        nu += Nd->getDOF();
                                }
                                else{
                                        cumap.addNode(umap_master[c_sys_index][u_index_tmp]);
                                }
                        } // Child constraint for loop ends here
                        cumap.finalize();
                }//for(int i=0;i<nConstraints;++i){
                
                // Size of a constraint jacobian is row x col where
                // row = number of algebraic constraints
                // col = number of u's tied by the constraint
                CJac_[i].resize(m,nu);
                CJac_[i].Zeros();
                delq_[i].resize(nu);
                
                // RhS of the equation CJac * delta q = -Cq
                Cq_[i].resize(m);
                Cq_[i].Zeros();
        }
        // TBD
        // printDebugInformation();
}

///////////////////////////////////////////////////////

void CoordinateProjection::printDebugInformation()const{        
        cout<<"\nConstraint ids in the system. Coupled constraints are shown in the "
        <<"same square brackets \n\n";
        for(int i=0;i<ccouple_.getLength();++i){
                const List <int>& tmp = ccouple_[i];
                cout<<"[ ";
                for(int j=0;j<tmp.getLength();++j){                        
                        cout<<tmp[j];
                        if(j!=tmp.getLength()-1)
                                cout<<", ";
                }
                cout<<" ]\n";
        }
        cout<<"\n--------------------------------------------------------\n";
        
        for(int ii=0;ii<constraints_.getLength();++ii){
                
                const List <int>& Parent_path = constraints_[ii]->getParentPathList();
                const List <int>& Child_path = constraints_[ii]->getChildPathList();
                
                for(int i=0;i<Parent_path.getLength();++i)
                        cout<<Parent_path[i]<<"\t";
                cout<<"("<<constraints_[ii]->getParentSysId()<<")"<<endl;
                for(int i=0;i<Child_path.getLength();++i)
                        cout<<Child_path[i]<<"\t";
                cout<<"("<<constraints_[ii]->getChildSysId()<<")"<<"\n---------\n";
        }
        
        cout<<"\n--------------------------------------------------------\n";
        
        for(int i=0;i<constraints_.getLength();++i){
                
                const List <int>& pumap = constraints_[i]->getPUMap();
                const List <int>& cumap = constraints_[i]->getCUMap();
                
                for(int i1=0;i1<pumap.getLength();++i1)
                        cout<<pumap[i1]<<"\t";
                cout<<"("<<constraints_[i]->getParentSysId()<<")"<<endl;
                for(int i1=0;i1<cumap.getLength();++i1)
                        cout<<cumap[i1]<<"\t";
                cout<<"("<<constraints_[i]->getChildSysId()<<")"<<"\n---------\n";
        }
        cout<<"\n--------------------------------------------------------\n";
        cout<<"Dimensions of the constraint Jacobians in use are\n";
        for(int i=0;i<ccouple_.getLength();++i)
                (CJac_+i)->printDimensions();
        
        cout<<"\n--------------------------------------------------------\n";
        
        for(int i=0;i<ConstraintToUMap_.getLength();++i){
                cout<<"Constraint "<<i << " : "<<ConstraintToUMap_[i]<<endl;
        }
}

/////////////////////////////////////////////////////// 